import os
import requests
import whois
from docx import Document
from fpdf import FPDF
import datetime

def investigate_website(url):
    """Investigate the website to find the country of operation and past security breaches."""
    try:
        # Perform WHOIS lookup to determine country
        domain = url.replace("https://", "").replace("http://", "").split("/")[0]
        whois_data = whois.whois(domain)
        country = whois_data.get("country", "Unknown")

        # Check for breaches using Have I Been Pwned API (requires API key)
        hibp_api_url = f"https://haveibeenpwned.com/api/v3/breachedaccount/{domain}"
        headers = {"hibp-api-key": "YOUR_API_KEY_HERE"}  # Replace with your HIBP API key
        breaches = []

        try:
            breach_response = requests.get(hibp_api_url, headers=headers)
            if breach_response.status_code == 200:
                breaches = breach_response.json()
            elif breach_response.status_code == 404:
                breaches = "No breaches found."
            else:
                breaches = f"Error: {breach_response.status_code}"
        except Exception as e:
            breaches = f"Error fetching breaches: {e}"

        return {
            "country": country,
            "breaches": breaches,
            "whois_data": str(whois_data)
        }
    except Exception as e:
        return {
            "country": "Error determining country",
            "breaches": f"Error fetching breaches: {e}",
            "whois_data": f"Error: {e}"
        }

def analyze_headers(url):
    """Perform a header analysis to identify security weaknesses."""
    try:
        response = requests.get(url)
        headers = response.headers

        # Security headers to check
        security_headers = {
            "Content-Security-Policy": "Missing or insufficient Content-Security-Policy",
            "Strict-Transport-Security": "Missing Strict-Transport-Security",
            "X-Content-Type-Options": "Missing X-Content-Type-Options",
            "X-Frame-Options": "Missing X-Frame-Options",
            "X-XSS-Protection": "Missing or insufficient X-XSS-Protection",
            "Referrer-Policy": "Missing Referrer-Policy",
            "Permissions-Policy": "Missing Permissions-Policy"
        }

        weaknesses = []
        for header, issue in security_headers.items():
            if header not in headers:
                weaknesses.append(issue)

        grade = "A" if not weaknesses else "C" if len(weaknesses) < 3 else "F"

        return {
            "grade": grade,
            "weaknesses": weaknesses,
            "headers": dict(headers)
        }
    except Exception as e:
        return {
            "grade": "Error",
            "weaknesses": [f"Error analyzing headers: {e}"],
            "headers": f"Error: {e}"
        }

def save_results_to_files(data, output_dir, domain):
    """Save investigation results to docx, pdf, and markdown files."""
    os.makedirs(output_dir, exist_ok=True)

    # Generate unique filenames using timestamp and domain
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    # Ensure the domain is a valid filename
    domain_safe = "".join(c if c.isalnum() else "_" for c in domain)
    base_filename = f"{domain_safe}_{timestamp}"
    docx_path = os.path.join(output_dir, f"{base_filename}.docx")
    pdf_path = os.path.join(output_dir, f"{base_filename}.pdf")
    md_path = os.path.join(output_dir, f"{base_filename}.md")

    # Create a Word document
    doc = Document()
    doc.add_heading("Website Security Analysis", level=1)

    doc.add_heading("1. Website Investigation", level=2)
    doc.add_paragraph(f"Country of Operation: {data['investigation'].get('country', 'N/A')}")
    doc.add_paragraph(f"Security Breaches: {data['investigation'].get('breaches', 'N/A')}")

    doc.add_heading("2. Header Analysis", level=2)
    doc.add_paragraph(f"Security Grade: {data['headers'].get('grade', 'N/A')}")
    if data['headers'].get('weaknesses'):
        doc.add_paragraph("Weaknesses:")
        for weakness in data['headers']['weaknesses']:
            doc.add_paragraph(f"- {weakness}", style="List Bullet")

    # Add Appendix
    doc.add_heading("Appendix", level=1)
    doc.add_heading("Full WHOIS Data", level=2)
    doc.add_paragraph(data['investigation'].get('whois_data', 'N/A'))

    doc.add_heading("Full Header Analysis", level=2)
    doc.add_paragraph(str(data['headers'].get('headers', 'N/A')))

    doc.add_heading("Full Breach Data", level=2)
    doc.add_paragraph(str(data['investigation'].get('breaches', 'N/A')))

    doc.save(docx_path)

    # Create a Markdown file
    with open(md_path, "w") as md_file:
        md_file.write("# Website Security Analysis\n\n")
        md_file.write("## 1. Website Investigation\n")
        md_file.write(f"- Country of Operation: {data['investigation'].get('country', 'N/A')}\n")
        md_file.write(f"- Security Breaches: {data['investigation'].get('breaches', 'N/A')}\n\n")
        md_file.write("## 2. Header Analysis\n")
        md_file.write(f"- Security Grade: {data['headers'].get('grade', 'N/A')}\n")
        if data['headers'].get('weaknesses'):
            md_file.write("- Weaknesses:\n")
            for weakness in data['headers']['weaknesses']:
                md_file.write(f"  - {weakness}\n")

        md_file.write("\n# Appendix\n\n")
        md_file.write("## Full WHOIS Data\n")
        md_file.write(f"```\n{data['investigation'].get('whois_data', 'N/A')}\n```\n")

        md_file.write("## Full Header Analysis\n")
        md_file.write(f"```\n{data['headers'].get('headers', 'N/A')}\n```\n")

        md_file.write("## Full Breach Data\n")
        md_file.write(f"```\n{data['investigation'].get('breaches', 'N/A')}\n```\n")

    # Create a PDF file using FPDF
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()
    pdf.set_font("Arial", size=12)

    pdf.cell(200, 10, txt="Website Security Analysis", ln=True, align='C')

    pdf.ln(10)
    pdf.set_font("Arial", style="B", size=12)
    pdf.cell(0, 10, "1. Website Investigation", ln=True)

    pdf.set_font("Arial", size=12)
    pdf.multi_cell(0, 10, f"Country of Operation: {data['investigation'].get('country', 'N/A')}")
    pdf.multi_cell(0, 10, f"Security Breaches: {data['investigation'].get('breaches', 'N/A')}")

    pdf.ln(5)
    pdf.set_font("Arial", style="B", size=12)
    pdf.cell(0, 10, "2. Header Analysis", ln=True)

    pdf.set_font("Arial", size=12)
    pdf.multi_cell(0, 10, f"Security Grade: {data['headers'].get('grade', 'N/A')}")
    if data['headers'].get('weaknesses'):
        pdf.cell(0, 10, "Weaknesses:", ln=True)
        for weakness in data['headers']['weaknesses']:
            pdf.cell(0, 10, f"- {weakness}", ln=True)

    pdf.add_page()
    pdf.set_font("Arial", style="B", size=12)
    pdf.cell(0, 10, "Appendix", ln=True)

    pdf.set_font("Arial", style="B", size=12)
    pdf.cell(0, 10, "Full WHOIS Data", ln=True)
    pdf.set_font("Arial", size=10)
    pdf.multi_cell(0, 10, data['investigation'].get('whois_data', 'N/A'))

    pdf.set_font("Arial", style="B", size=12)
    pdf.cell(0, 10, "Full Header Analysis", ln=True)
    pdf.set_font("Arial", size=10)
    pdf.multi_cell(0, 10, str(data['headers'].get('headers', 'N/A')))

    pdf.set_font("Arial", style="B", size=12)
    pdf.cell(0, 10, "Full Breach Data", ln=True)
    pdf.set_font("Arial", size=10)
    pdf.multi_cell(0, 10, str(data['investigation'].get('breaches', 'N/A')))

    pdf.output(pdf_path)

    print(f"Results saved to: {docx_path}, {md_path}, {pdf_path}")

if __name__ == "__main__":
    # User inputs
    website_url = input("Enter the website URL (including https://): ").strip()
    default_directory = os.getcwd()
    print(f"Default save directory: {default_directory}")
    use_default = input("Do you want to use the default directory? (y/n): ").strip().lower()

    if use_default == 'y':
        output_directory = default_directory
    else:
        output_directory = input("Enter the directory to save results: ").strip()

    # Extract domain name for unique file naming
    domain = website_url.replace("https://", "").replace("http://", "").split("/")[0]

    # Perform analysis
    investigation_results = investigate_website(website_url)
    header_analysis_results = analyze_headers(website_url)

    # Combine results
    results = {
        "investigation": investigation_results,
        "headers": header_analysis_results
    }

    # Save to files
    save_results_to_files(results, output_directory, domain)
